#include <windows.h>
#include <Userenv.h>
#include <tchar.h>
#include <stdio.h>
#include <Tlhelp32.h>

WINAPI DWORD WTSGetActiveConsoleSessionId();
WINAPI BOOL WTSQueryUserToken(ULONG SessionId,PHANDLE phToken);
WINAPI BOOL ProcessIdToSessionId(DWORD dwProcessId,DWORD *pSessionId);

DWORD getWinlogonPid(DWORD session){
    PROCESSENTRY32 pe;
    pe.dwSize = sizeof(pe);

    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,0);

    BOOL good = Process32First(snap,&pe);
    while(good){
        if (strcmp(pe.szExeFile,"winlogon.exe")==0){
            DWORD processsession = 0;
            if (ProcessIdToSessionId(pe.th32ProcessID,&processsession)){
                if (processsession == session){
                    return pe.th32ProcessID;
                }
            }
        }
        good = Process32Next(snap,&pe);
    }
    return 0;
}


BOOL runOnScreen(const char *executable){
    SetProcessWindowStation(OpenWindowStation("WinSta0",FALSE,GENERIC_ALL));

    DWORD session = WTSGetActiveConsoleSessionId();
    if (session == 0xffffffff) return FALSE;

    DWORD winlogonpid = getWinlogonPid(session);
    if (!winlogonpid) return FALSE;

    HANDLE winlogonhandle = OpenProcess(MAXIMUM_ALLOWED,FALSE,winlogonpid);
    if (!winlogonhandle) return FALSE;

    HANDLE hProcessToken;
    if (!OpenProcessToken(winlogonhandle,TOKEN_ADJUST_PRIVILEGES|TOKEN_QUERY|TOKEN_DUPLICATE|TOKEN_ASSIGN_PRIMARY|TOKEN_READ|TOKEN_WRITE,&hProcessToken))return FALSE;

    LUID luid;
    if (!LookupPrivilegeValue(NULL,SE_DEBUG_NAME,&luid)) return FALSE;

    TOKEN_PRIVILEGES tp;
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    HANDLE hTokenDup;
    if (!DuplicateTokenEx(hProcessToken,MAXIMUM_ALLOWED,NULL,SecurityIdentification,TokenPrimary, &hTokenDup)) return FALSE;

    SetTokenInformation(hTokenDup,TokenSessionId,(void*)session,sizeof(DWORD));

    if (!AdjustTokenPrivileges(hTokenDup,FALSE,&tp,sizeof(TOKEN_PRIVILEGES),(PTOKEN_PRIVILEGES)NULL,NULL)) return FALSE;

    void *lpEnv;
    if (!CreateEnvironmentBlock(&lpEnv, hTokenDup, TRUE)) return FALSE;

    HDESK hDesktop = OpenInputDesktop(0,0,0);
    if (hDesktop == NULL) return FALSE;

    TCHAR desktopName[256];
    DWORD needed;
    if (!GetUserObjectInformation(hDesktop,UOI_NAME,&desktopName,256,&needed)) return FALSE;
    desktopName[needed]='\0';

    CloseDesktop(hDesktop);

    STARTUPINFO si;
    PROCESS_INFORMATION pi;
    memset(&si, 0, sizeof(si));
    memset(&pi, 0, sizeof(pi));
    si.cb = sizeof(si);
    char desktop[strlen(desktopName)+9];
    strcpy(desktop,"winsta0\\");
    strcat(desktop,desktopName);
    si.lpDesktop = desktop;

    char command[strlen(executable)+1];
    strcpy(command,executable);

    BOOL good = CreateProcessAsUser(hTokenDup,NULL,command,NULL,NULL,FALSE,CREATE_UNICODE_ENVIRONMENT,lpEnv,NULL,&si,&pi);

    CloseHandle(hProcessToken);
    CloseHandle(hTokenDup);

    return good;
}
